// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: jhkim

#include <algorithm>
#include "coin2/strategy/linear_model/focus2.h"
#include <numeric>

std::vector<double> calculate_average(const std::vector<std::vector<double>>& input) {
  if (input.empty() || input[0].empty()) return {};
  std::vector<double> average_vector(input[0].size(), 0.0);
  for (const auto& row : input) {
    std::transform(row.begin(), row.end(), average_vector.begin(), average_vector.begin(),
                   [](double value, double accumulator) { return accumulator + value; });
  }
  for (double& val : average_vector) {
    val /= input.size();
  }
  return average_vector;
}

namespace coin2::strategy::linear_model {

void LmSubStrategy::OnBufferSignal(
    int64_t timestamp,
    int64_t* first_feed_time,
    double midret_pred,
    double signal_sigma,
    OrderExecutorSystem* oe,
    PredictionInterval* pred_interval,
    AggregatedSizer* asizer) {
  if (*first_feed_time == 0L) *first_feed_time = timestamp;
  if (timestamp - *first_feed_time < 30e9) return;
  double min_quote = GetMinPosQuote(focus_config) * focus_config.buffer_max_pos_multiplier();
  double max_quote = GetMaxPosQuote(focus_config) * focus_config.buffer_max_pos_multiplier();
  double reserve_quote = GetQuoteFromPq(*pholder, bbo->askp, reserve);

  bool print_log = false;
  if (last_logged_ts_ == 0L) last_logged_ts_ = timestamp;
  if (last_logged_ts_ + 120'000'000'000L < timestamp) {
    print_log = true;
    last_logged_ts_ = timestamp;
  }

  if (print_log) {
    LOG(INFO) << symbol->symbol_string() << ": signal=" << midret_pred;
  }

  double curr_pos = GetCurrPos(oe);
  double curr_pos_quote = GetQuoteFromPq(*pholder, bbo->askp, curr_pos);
  double desired_pos_quote = asizer->DesiredBufferPosQuote();

  if (asizer->symbols_clear.count(symbol) > 0) {
    min_quote = reserve_quote;
    max_quote = reserve_quote;
    desired_pos_quote = reserve_quote;
  }
  double max_pos = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      max_quote);
  double min_pos = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      min_quote);
  double lot_size = (max_pos - min_pos) / focus_config.stack();
  double min_order_qty = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      focus_config.min_order_quote());

  double pos_quote_diff = 0;
  double signal = 0;
  if (curr_pos_quote < desired_pos_quote) {
    pos_quote_diff = desired_pos_quote - curr_pos_quote;
    signal = midret_pred + focus_config.buffer_skew_bps() * pos_quote_diff / (max_quote - min_quote);
    if (!focus_config.buffer_allow_divergence()) {
      signal = (signal > 0)? signal : 0;
    }
  } else if (curr_pos_quote > desired_pos_quote) {
    pos_quote_diff = curr_pos_quote - desired_pos_quote;
    signal = midret_pred - focus_config.buffer_skew_bps() * pos_quote_diff / (max_quote - min_quote);
    if (!focus_config.buffer_allow_divergence()) {
      signal = (signal < 0)? signal : 0;
    }
  }

  if (print_log) {
    LOG(INFO) << symbol->symbol_string() << ": adjusted_signal=" << signal;
  }

  lot_size = std::min(
      lot_size,
      util::GetQtyFromQuoteSafe(
          *pholder,
          price_window_.last_element(),
          price_window_.first_element(),
          pos_quote_diff));
  oe->AdjustSize(*pholder, min_pos, max_pos, lot_size, min_order_qty);

  ManageMmOrder4(
    oe,
    *pholder,
    *bbo,
    timestamp,

    &bom,
    leverage_rate,

    signal,
    signal,
    signal_sigma,

    std::nullopt,
    std::nullopt,

    false,
    false,

    focus_config);

  oe->onPeriodicOrderCheck(*pholder, timestamp);
}

void LmSubStrategy::UpdateSignal2(
    int64_t timestamp,
    int64_t* first_feed_time,
    double midret_pred,
    std::optional<double> midret_pred_sell,
    std::optional<double> midtomed_pred,
    std::optional<double> midtomax_pred,
    std::optional<double> midtomin_pred,
    std::optional<double> residual_pred,
    OrderExecutorSystem* oe,
    PredictionInterval* pred_interval,
    AggregatedSizer* asizer,
    std::vector<double> signals) {
  if (!bbo) return;
  if (!pholder) return;
  CHECK(oe);
  if (!oe->HasProduct(*pholder)) return;
  if (!oe->HasPosition(*pholder)) return;

  midret_pred *= focus_config.signal_multiplier();
  pred_interval->UpdateSignal(timestamp, midret_pred);

  double signal_sigma;
  if (focus_config.has_fixed_sigma()) {
    signal_sigma = focus_config.fixed_sigma();
  } else {
    double _mean;
    bool use_warmup = focus_config.has_lean_aggr_nsigma() || focus_config.has_lean_per_symbol_nsigma() || focus_config.has_buy_threshold_nsigma();
    if(!pred_interval->GetMeanStd(&_mean, &signal_sigma, use_warmup)) return;
  }

  if (focus_config.use_midtomed_signal()) {
    midret_pred = *midtomed_pred;
  }

  if (focus_config.use_residual_signal()) {
    CHECK(!focus_config.use_midtomed_signal());
    midret_pred += *residual_pred * focus_config.residual_signal_multiplier();
  }

  if (focus_config.buffer_mode() && is_buffer_symbol) {
    CHECK(focus_config.use_reserve() == 0);
    OnBufferSignal(timestamp, first_feed_time, midret_pred, signal_sigma, oe, pred_interval, asizer);
    return;
  }

  if (last_logged_ts_ == 0L) last_logged_ts_ = timestamp;
  if (last_logged_ts_ + 120'000'000'000L < timestamp) {
    if (focus_config.use_sell_model()) {
      LOG(INFO) << symbol->symbol_string() << ": buy_signal=" << midret_pred << ", sell_signal=" << *midret_pred_sell;
    } else {
      LOG(INFO) << symbol->symbol_string() << ": signal=" << midret_pred;
    }
    last_logged_ts_ = timestamp;
  }

  const auto& pi = pholder->product_info();
  double contract_value = pi.contract_value() > 0? pi.contract_value(): 1.0;

  double min_quote = GetMinPosQuote(focus_config, 1.0) * focus_config.max_concentration();
  double max_quote = GetMaxPosQuote(focus_config, 1.0) * focus_config.max_concentration();
  double reserve_quote = GetQuoteFromPq(*pholder, bbo->askp, reserve);

  if (asizer->symbols_clear.count(symbol) > 0) {
    min_quote = reserve_quote;
    max_quote = reserve_quote;
  }

  if (focus_config.liquidity_constraint() > 0) {
    liquidity_.UpdateTimeOnly(timestamp);
    double liquidity_sum = liquidity_.sum();
    if (pi.is_inverse()) {
      liquidity_sum /= bbo->askp;
    }
    double max_absquote = contract_value * liquidity_sum * focus_config.liquidity_constraint();
    min_quote = std::max(min_quote, -max_absquote);
    max_quote = std::min(max_quote, max_absquote);
  }

  double curr_pos = GetCurrPos(oe);

  trading_pnls_.UpdateTimeOnly(timestamp);
  holding_pnls_.UpdateTimeOnly(timestamp);

  double midp = 0.5 * (bbo->askp + bbo->bidp);

  double holding_pnl = curr_pos * midp;
  if (holding_pnls_.num_popped() > 0) {
    holding_pnl -= holding_pnls_.last_popped();
  }

  double pnl_q = 0;
  if (trading_pnls_.num_popped() > 0) {
    pnl_q = (trading_pnls_.sum() + holding_pnl) / midp;
  }

  double curr_pos_quote = GetQuoteFromPq(
      *pholder,
      bbo->askp,
      curr_pos);
  asizer->UpdatePosition(symbol, curr_pos_quote);
  auto bs = asizer->GetRemainingBuySellQuote();
  double aggr_norm_pos = asizer->GetAggregatedNormalizedPos(focus_config.use_reserve());
  double aggr_skew_bps;
  if (focus_config.has_lean_aggr_nsigma()) {
    aggr_skew_bps = -aggr_norm_pos * focus_config.lean_aggr_nsigma() * signal_sigma;
  } else {
    aggr_skew_bps = -aggr_norm_pos * focus_config.lean_aggr_bps();
  }

  double max_pos_config = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      max_quote);
  double min_pos_config = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      min_quote);
  double lot_size_config = (max_pos_config - min_pos_config) / focus_config.stack();
  double min_order_qty = util::GetQtyFromQuoteSafe(
      *pholder,
      price_window_.last_element(),
      price_window_.first_element(),
      focus_config.min_order_quote());

  double max_pos = max_pos_config;
  double min_pos = min_pos_config;
  double lot_size = lot_size_config;
  if (signals.size() >= 2) {
    // const double horiz = 900.;
    // double uncertainty = signals[1];
    // TODO: parametrize or get from horiz
    // double w_inv_norm = horiz / uncertainty;
    // lot_size *= std::min(w_inv_norm, 1.);

    if (signals[1] > 0) {
      double uncertainty = signals[1];
      double adj = std::sqrt(uncertainty) * 0.2;
      if (midret_pred > 0) {
        midret_pred = std::max(0., midret_pred - adj);
      } else {
        midret_pred = std::min(0., midret_pred + adj);
      }
    }
  }

  if (focus_config.use_aggregated_pos()) {
    double remain_buy_qty = std::max(0., GetQtyFromQuote(*pholder, bbo->askp, bs.first));
    double remain_sell_qty = std::max(0., GetQtyFromQuote(*pholder, bbo->askp, bs.second));
    max_pos = std::min(curr_pos + remain_buy_qty, max_pos);
    min_pos = std::max(curr_pos - remain_sell_qty, min_pos);
  }

  bool buy_block = false, sell_block = false;
  // limit exposure
  if (focus_config.use_aggregated_pos()) {
    buy_block |= bs.first <= 0;
    sell_block |= bs.second <= 0;
  }

  // limit abspos quote
  if (focus_config.max_abspos_quote() > 0) {
    auto bs2 = asizer->GetRemainingAbsBuySellQuote(symbol);
    buy_block |= bs2.first <= 0;
    sell_block |= bs2.second <= 0;
  }

  if (focus_config.fill_cooldown_constraint() > 0) {
    buy_filled_.UpdateTimeOnly(timestamp);
    sell_filled_.UpdateTimeOnly(timestamp);
    buy_block |= buy_filled_.sum() > lot_size * focus_config.fill_cooldown_constraint();
    sell_block |= sell_filled_.sum() > lot_size * focus_config.fill_cooldown_constraint();
  }

  double cur_stack = GetCurrStack(oe, lot_size);
  if (focus_config.fill_cooltime_sec2() > 0) {
    if (cur_stack > 0) {
      int64_t cool_until = bom.prev_bid_fill_time + focus_config.fill_cooltime_sec2() * 1e9 * cur_stack / (focus_config.stack() * 0.5);
      buy_block = buy_block || (timestamp < cool_until);
    } else if (cur_stack < 0) {
      int64_t cool_until = bom.prev_ask_fill_time + focus_config.fill_cooltime_sec2() * 1e9 * (-cur_stack) / (focus_config.stack() * 0.5);
      sell_block = sell_block || (timestamp < cool_until);
    }
  }

  if (focus_config.pr_constraint() > 0) {
    buy_turnover_.UpdateTimeOnly(timestamp);
    sell_turnover_.UpdateTimeOnly(timestamp);
    market_turnover_.UpdateTimeOnly(timestamp);
    double max_pr = focus_config.pr_constraint();
    double market_turnover = market_turnover_.sum();
    double turnover = buy_turnover_.sum() + sell_turnover_.sum();
    bool pr_exceed = market_turnover * max_pr < turnover;
    buy_block |= pr_exceed;
    sell_block |= pr_exceed;
  }

  if (focus_config.book_constraint() > 0) {
    double ratio = focus_config.book_constraint();

    ask_bbo_.UpdateTimeOnly(timestamp);
    bid_bbo_.UpdateTimeOnly(timestamp);
    double ask_bbo = ask_bbo_.average();
    double bid_bbo = bid_bbo_.average();
    auto& tickprc = pholder->product_info().tick_price();
    double book_mult = tickprc.Infer2bpBookMult(bbo->askp);

    double remain_buy_qty = std::max(0.0, bid_bbo * book_mult * ratio / bbo->askp);
    double remain_sell_qty = std::max(0.0, ask_bbo * book_mult * ratio / bbo->askp);

    max_pos = std::min(max_pos, curr_pos + remain_buy_qty);
    min_pos = std::max(min_pos, curr_pos - remain_sell_qty);
  }

  if (pnl_q < 0) {
    min_pos -= focus_config.loss_penalty_ratio() * pnl_q;
    max_pos += focus_config.loss_penalty_ratio() * pnl_q;
  }

  oe->AdjustSize(*pholder, min_pos, max_pos, lot_size, min_order_qty);

  if (focus_config.use_sell_model()) {
    *midret_pred_sell *= focus_config.signal_multiplier();
  }

  if (focus_config.use_zscore_signal()) {
    double ret_mean, ret_std;
    if (!pred_interval->GetMeanStd(&ret_mean, &ret_std, focus_config.zscore_signal_warmup()) || ret_std <= 0.0) return;
    midret_pred = (midret_pred - ret_mean) / ret_std;
  }

  double rate_used_ratio = 0;
  auto* rate_tracker = oe->order_system()->gateway()->GetMutableRateTracker();
  if (rate_tracker && focus_config.enable_api_rate_limit_check()) {
    rate_used_ratio = std::max(
        rate_tracker->GetMaxRateUsedRatio(timestamp, ActionType::PLACE),
        rate_tracker->GetMaxRateUsedRatio(timestamp, ActionType::CANCEL));
  }

  if (*first_feed_time == 0L) *first_feed_time = timestamp;
  if (timestamp - *first_feed_time < 30e9) return;
  double current_rel_pos = focus_config.use_reserve() ? (curr_pos - 0.5 * (min_pos_config + max_pos_config)) : curr_pos;
  double current_stack = current_rel_pos / lot_size;
  double skew_bps;
  if (focus_config.has_lean_per_symbol_nsigma()) {
    skew_bps = -current_stack / focus_config.stack() * 2 * focus_config.lean_per_symbol_nsigma() * signal_sigma;
  } else {
    skew_bps = -current_stack / focus_config.stack() * 2 * focus_config.lean_per_symbol_bps();
  }
  if (focus_config.enable_api_rate_limit_check() &&
      rate_used_ratio > focus_config.reduce_order_above_api_rate() &&
      focus_config.reduce_order_multiplier() > 0) {
    skew_bps *= focus_config.reduce_order_multiplier() / (1 - rate_used_ratio + 0.01);
  }
  double signal_skewed_lb = midret_pred;
  double signal_skewed_ub = (focus_config.use_sell_model())? *midret_pred_sell : midret_pred;
  if (skew_bps + aggr_skew_bps > 0) {
    // aggressive bid posting becomes less aggressive
    signal_skewed_lb += (skew_bps + aggr_skew_bps) * focus_config.lean_close_multiplier();
    signal_skewed_ub += skew_bps + aggr_skew_bps;
  } else {
    // aggressive ask posting becomes less aggressive
    signal_skewed_lb += skew_bps + aggr_skew_bps;
    signal_skewed_ub += (skew_bps + aggr_skew_bps) * focus_config.lean_close_multiplier();
  }

  if ((!focus_config.enable_api_rate_limit_check() || rate_used_ratio < focus_config.layering_until_api_rate()) &&
      focus_config.layering() > 1) {
    ManageMmOrderLayering(
        oe,
        *pholder,
        *bbo,
        timestamp,

        &bom,
        leverage_rate,

        signal_skewed_lb,
        signal_skewed_ub,
        signal_sigma,

        (focus_config.no_block_on_layering() > 0)? false : buy_block,
        (focus_config.no_block_on_layering() > 0)? false : sell_block,

        focus_config,

        rate_used_ratio);
  } else {
    ManageMmOrder4(
        oe,
        *pholder,
        *bbo,
        timestamp,

        &bom,
        leverage_rate,

        signal_skewed_lb,
        signal_skewed_ub,
        signal_sigma,

        midtomax_pred,
        midtomin_pred,

        buy_block,
        sell_block,

        focus_config,

        rate_used_ratio);
  }
  oe->onPeriodicOrderCheck(*pholder, timestamp);
}


void LmFocusLogic2::onFeed(
    const Symbol* symbol,
    const FeedUpdate& upd,
    const ::fastfeature::FeedMessage& feed_msg,
    LinearModel* linear_model,
    int64_t* lm_feed_id,
    const int64_t* global_feed_id,
    LinearModelDelegateProto* delegate,
    bool simulation) {
  if (strats_.count(symbol) > 0) {
    strats_.at(symbol)->UpdateProduct(product_cache_, upd.product(), upd.timestamp());
    if (feed_msg.has_book()
        && feed_msg.book().ask_size() > 0
        && feed_msg.book().bid_size() > 0) {
      auto ask0 = feed_msg.book().ask(0);
      auto bid0 = feed_msg.book().bid(0);
      auto bbo = Bbo {
          feed_msg.fetched_time(),
          feed_msg.exchange_time(),
          ask0.price_float(),
          bid0.price_float(),
      };
      strats_.at(symbol)->UpdateBbo(bbo);

      double ask0_pq = ask0.price_float() * ask0.qty_float();
      double bid0_pq = bid0.price_float() * bid0.qty_float();
      strats_.at(symbol)->UpdateBboPq(ask0_pq, bid0_pq, feed_msg.fetched_time());

      strats_.at(symbol)->ManageExecutorConfig(upd.timestamp());
    }
    if (feed_msg.feed_type() == ::fastfeature::TRADE_FEED)
      strats_.at(symbol)->UpdateTrade(
          feed_msg.trade().price_float(),
          feed_msg.trade().qty_float(),
          feed_msg.fetched_time());
  }

  // LOG(INFO) << linear_model->GetAllFeatures()->DebugString();
  if (asizer_.symbols_clear.count(symbol) > 0) {
    CHECK(strats_.count(symbol) > 0);
    strats_[symbol]->ClearPositionWoSignal(
        upd.timestamp(),
        oe_,
        &asizer_);
    return;
  }
  delegate->ClearProto();
  lholder_.TriggerModelWithLatencyRecording(linear_model, nullptr);

  if (delegate->response_.buy_model().empty()) {
    if (strats_.count(symbol) > 0) {
      if (last_logged_tss_[symbol] == 0L) last_logged_tss_[symbol] = feed_msg.fetched_time();
      if (last_logged_tss_[symbol] + 120'000'000'000L < feed_msg.fetched_time()) {
        LOG(INFO) << symbol->symbol_string() << ": buy_model is empty";
        last_logged_tss_[symbol] = feed_msg.fetched_time();
      }
    }
    return;
  }
  auto signal_resp = delegate->response_.buy_model().Get(0);
  double buymodel_signal = 0.0, sellmodel_signal = 0.0;
  int n_buy_models = 0, n_sell_models = 0, n_residual_models = 0;
  std::vector<std::vector<double>> signals_models;
  std::vector<double> signals;
  for (int i = 0; i < delegate->response_.buy_model().size(); i++) {
    const auto& buy_model = delegate->response_.buy_model().Get(i);
    if (!lm_product_config_.focus_config().ignore_is_model() ||
        IsOosModel(upd.timestamp(), buy_model)) {
      if (lm_product_config_.focus_config().signal_multipliers().size() > 0) {
        CHECK_EQ(
            delegate->response_.buy_model().size(),
            lm_product_config_.focus_config().signal_multipliers().size());
        buymodel_signal += buy_model.signal() * lm_product_config_.focus_config().signal_multipliers()[i];
        signals_models.push_back({buy_model.signal_multi().begin(), buy_model.signal_multi().end()});
      } else {
        buymodel_signal += buy_model.signal();
        signals_models.push_back({buy_model.signal_multi().begin(), buy_model.signal_multi().end()});
      }
      ++n_buy_models;
    }
  }
  if (n_buy_models > 0) {
    buymodel_signal /= n_buy_models;
    signals = calculate_average(signals_models);
  }

  // use_nth_model_only begins from 1; if zero, ensemble as above (default).
  int32_t model_idx = lm_product_config_.focus_config().use_nth_model_only();
  if (model_idx > 0) {
    CHECK_LE(model_idx, delegate->response_.buy_model().size());
    buymodel_signal = delegate->response_.buy_model().Get(model_idx - 1).signal();
    signals = signals_models[model_idx - 1];
  }

  if (lm_product_config_.focus_config().use_sell_model()) {
    for (int i = 0; i < delegate->response_.sell_model().size(); i++) {
      const auto& sell_model = delegate->response_.sell_model().Get(i);
      if (!lm_product_config_.focus_config().ignore_is_model() ||
          IsOosModel(upd.timestamp(), sell_model)) {
        sellmodel_signal += sell_model.signal();
        ++n_sell_models;
      }
    }
    if (n_sell_models > 0) sellmodel_signal /= n_sell_models;
  }

  std::optional<double> midtomed_pred;
  if (!delegate->response_.midtomed_model().empty()) {
    midtomed_pred = delegate->response_.midtomed_model().Get(0).signal();
  }

  std::optional<double> midtomax_pred;
  if (!delegate->response_.midtomax_model().empty()) {
    midtomax_pred = delegate->response_.midtomax_model().Get(0).signal();
  }

  std::optional<double> midtomin_pred;
  if (!delegate->response_.midtomin_model().empty()) {
    midtomin_pred = delegate->response_.midtomin_model().Get(0).signal();
  }

  std::optional<double> residual_pred;
  if (!delegate->response_.residual_model().empty()) {
    double resmodel_signal = 0.0;
    for (int i = 0; i < delegate->response_.residual_model().size(); i++) {
      const auto& residual_model = delegate->response_.residual_model().Get(i);
      if (!lm_product_config_.focus_config().ignore_is_model() ||
          IsOosModel(upd.timestamp(), residual_model)) {
        resmodel_signal += residual_model.signal();
        ++n_residual_models;
      }
    }
    if (n_residual_models > 0) resmodel_signal /= n_residual_models;
    residual_pred = resmodel_signal;
  }

  auto focus_symbolstr = signal_resp.symbol();
  const Symbol* focus_symbol = linear_model->GetFeatureManager()
      ->GetSymbolManager()->FromSymbolString(focus_symbolstr);

  if (asizer_.symbols_clear.count(focus_symbol) == 0 && strats_.count(focus_symbol) > 0) {
    strats_[focus_symbol]->UpdateSignal2(
        upd.timestamp(),
        &first_feed_time_,
        buymodel_signal,
        sellmodel_signal,
        midtomed_pred,
        midtomax_pred,
        midtomin_pred,
        residual_pred,
        oe_,
        &pred_interval_,
        &asizer_,
        signals);
  }
  if (!simulation) {
    lholder_.LogSignal(upd.timestamp(), delegate);
  }
}

}  // namespace coin2::strategy::linear_model
